############################################################################################
# Read the simulation results and preprocess them, then save to .bin file for later usage
#
#


import glob
import os
import time
from pathlib import Path

import numpy as np
import pandas as pd

# =========================================================
# Normliazation functions:
#
# =========================================================


def normalize_by_max(arr):
    return arr / np.max(arr)


def normalize_by_sum(arr, binwidth):
    sum = np.sum(arr)
    if sum == 0:
        return np.zeros_like(arr)
    return arr / sum / binwidth


# =========================================================
# File path info
# - parname: H, He, Li, Be, B, C, O, Ne
# - group: 代表模拟结果分类存储的信息: c, c_nonprim, h, he, li, be...., other, any等等
# =========================================================

particle_name = "he"
engy_dir = "275MeV"
dir_tag = "filter"
# group = 'c'

# path = '/data/wenxiao/cmcvalidate/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/'
# path = '/data/wenxiao/cmcvalidate/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/eDep_filter_' + group + '.txt'
# path = '/data/wenxiao/cmcvalidate/cudaCMC_validate_Geant4/G4Project/C_400MeV_filter/eDep_filter_c.txt'

# path = '/data/wenxiao/cmcvalidate_1/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/'
# path = '/data/wenxiao/cmcvalidate_2/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/'
# path = '/data/wenxiao/cmcvalidate_3/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/'

# path = '/data/wenxiao/cmcvalidate_formodel/cudaCMC_validate_Geant4/G4Project/' + parname + '_' + enedir + '_' + dirtag + '/'

path = "/data/wenxiao/cmcvalidate_letd_new/cudaCMC_validate_Geant4/G4Project/" + particle_name + "_" + engy_dir + "/"
# arr0 = np.loadtxt(path, delimiter=',')

txt_files = glob.glob(os.path.join(path, "eDep.txt"))
print(txt_files)


# =========================================================
# Simulation: num of particles
# =========================================================
num_of_simult_prtcls = 1e6


# =========================================================
# Reshape setting
# - Keep size the same with simulation mesh box setting
# =========================================================

Zbin = 1000
Zmax = 50.0
Ybin = 1000
Ymax = 50
Ymin = -50

# Zaxis = np.arange(Zmax / Zbin / 2., Zmax, Zmax / Zbin)
# print('Zaxis: {}'.format(Zaxis))
# Yaxis = np.arange(Ymin + (Ymax - Ymin) / Ybin / 2, Ymax, (Ymax - Ymin) / Ybin)

# =========================================================
## 数据一共六列,
## idx0, idx1, idx2, value, val2, entry

for txt_file in txt_files:
    start_time = time.time()
    try:
        print(f"Processing file: {txt_file}")
        # arr0 = np.loadtxt(txt_file, delimiter=',')
        arr0 = pd.read_csv(txt_file, delimiter=",", comment="#", header=None).values  # DataFrame 转为numpy array

        arr0 = arr0.reshape(Ybin, Ybin, Zbin, 6)
        print(arr0.shape)
        ## get data of total value
        all0 = arr0[:, :, :, 3]
        # all0 = arr0
        print(all0.shape)
        # print(all0)
        ## 交换次序 -> Z X Y
        all0 = np.transpose(all0, (2, 0, 1))
        # print(all0)
        totDose0 = all0[:, :, :] / num_of_simult_prtcls

        doseY0 = np.sum(totDose0, axis=(0, 2)) # Z 求和; Y 求和
        doseZ0 = np.sum(totDose0, axis=(1, 2)) # X 求和; Y 求和
        # norm_doseZ0 = normalize_by_max(doseZ0)
        # norm_doseZ0 = normalize_by_sum(doseZ0, Zmax / Zbin)
        doseY0_Brag_raw = np.sum(totDose0, axis=2)

        doseY0_Brag = np.empty((Zbin, Ybin))
        # 进行归一化
        #        doseY0_Brag = np.apply_along_axis(normalize_by_sum, axis=1, arr=doseY0_Brag_raw, binwidth=(Ymax - Ymin) / Ybin)
        # for j in range(Zbin):
        #    doseY0_Brag[j] = normalize_by_sum(doseY0_Brag_raw[j])

        # 将处理后的doseY0_Brag存到bin文件中
        # bin_file = os.path.basename(path).replace('.txt', '.bin')
        # bin_file = Path(path).stem + '.bin'
        bin_file = txt_file.replace(".txt", ".bin")

        # Test: 直接存储未归一化的数据
        # doseY0_Brag = doseY0_Brag.astype(np.float64)
        doseY0_Brag = doseY0_Brag_raw.astype(np.float64)
        doseY0_Brag.tofile(bin_file)

        elapsed_time = time.time() - start_time
        print(f"Successfully saved to: {bin_file}")
        print(f"Processing time: {elapsed_time:.2f} seconds")

    except Exception as e:
        print(f"Error processing file {txt_file}: {str(e)}")
